'''
麻雀搜索算法
'''
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
import random
import os
import copy
from scipy.stats import norm
import time


class SSA2020():
    def __init__(self, n_dim=None, pop_size=20, max_iter=50, lb=-512, ub=512, obj_func=None, ax=None):
        self.dim = n_dim
        self.pop = pop_size
        self.ub = ub
        self.lb = lb
        self.Max_iter = max_iter
        self.fun = obj_func
        self.ax = ax  # chart对象，用于3d绘图描点
        self.X = np.zeros([pop_size, n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i, j] = np.random.random() * (ub[j] - lb[j]) + lb[j]
        if (self.ax):
            # 初始化标点
            self.ax.chat3dplot(self.X, self.fun, 'ax1', 'o', 'black')

    '''边界检查函数'''

    def BorderCheck(self):
        for i in range(self.pop):
            for j in range(self.dim):
                if self.X[i, j] > self.ub[j]:
                    self.X[i, j] = self.ub[j]
                elif self.X[i, j] < self.lb[j]:
                    self.X[i, j] = self.lb[j]

    '''计算适应度函数'''

    def CaculateFitness(self):
        pop = self.X.shape[0]
        fitness = np.zeros([pop, 1])
        for i in range(pop):
            fitness[i] = self.fun(self.X[i, :])
        return fitness

    '''适应度排序'''

    def SortFitness(self, Fit):
        fitness = np.sort(Fit, axis=0)
        index = np.argsort(Fit, axis=0)
        return fitness, index

    '''根据适应度对位置进行排序'''

    def SortPosition(self, index):
        Xnew = np.zeros(self.X.shape)
        for i in range(self.X.shape[0]):
            Xnew[i, :] = self.X[index[i], :]
        return Xnew

    '''麻雀搜索算法'''

    def Tent_SSA(self,isBox = False):
        pop = self.pop
        dim = self.dim
        Max_iter = self.Max_iter
        avg_fit = []  # 每次迭代的所有个体适应度平均值
        avg_ts = []  # 平均方差
        ST = 0.8
        PD = 0.2
        SD = 0.1
        PDNumber = int(pop * PD)
        SDNumber = int(pop * SD)
        fitness = self.CaculateFitness()
        fitness,sortIndex = self.SortFitness(fitness)
        self.X = self.SortPosition(sortIndex)
        GbestScore = copy.copy(fitness[0])
        GbestPosition = np.zeros([1,dim])
        GbestPosition[0,:] = copy.copy(self.X[0,:])
        Curve = np.zeros([Max_iter+1,1])
        Curve[0] = fitness[0]
        time_start = time.time()  # 记录迭代寻优开始时间
        for t in range(Max_iter):
            BestF = copy.copy(fitness[0])
            Xworst = copy.copy(self.X[-1,:])
            Xbest = copy.copy(self.X[0,:])
            '''发现者位置更新'''
            R2 = np.random.random()
            for i in range(PDNumber):
                if R2 < ST:
                    self.X[i,:] = self.X[i,:] * np.exp(-i/(np.random.random()*Max_iter))
                else:
                    self.X[i,:] = self.X[i,:] + np.random.randn()*np.ones([1,dim])
            self.BorderCheck()
            fitness = self.CaculateFitness()

            bestII = np.argmin(fitness)
            Xbest = copy.copy(self.X[bestII,:])
            '''加入者更新'''
            for i in range(PDNumber+1,pop):
                if i > (pop - PDNumber)/2+PDNumber:
                    self.X[i,:] = np.random.randn() * np.exp((Xworst - self.X[i,:])/i ** 2)
                else:
                    A = np.ones([dim,1])
                    for a in range(dim):
                        if(np.random.random() > 0.5):
                            A[a] = -1
                    AA = np.dot(A,np.linalg.inv(np.dot(A.T,A)))
                    self.X[i,:] = self.X[0,:] + np.abs(self.X[i,:] - GbestPosition)*AA.T
            self.BorderCheck()
            fitness = self.CaculateFitness()
            '''危险预警更新'''
            Temp = range(pop)
            RandIndex = random.sample(Temp,pop)
            SDchooseIndex = RandIndex[0:SDNumber]
            for i in range(SDNumber):
                if fitness[SDchooseIndex[i]] > BestF:
                    self.X[SDchooseIndex[i],:] = Xbest + np.random.randn() * np.abs(self.X[SDchooseIndex[i],:] -  Xbest)
                elif fitness[SDchooseIndex[i]] == BestF:
                    K = 2 * np.random.random() - 1
                    self.X[SDchooseIndex[i],:] = self.X[SDchooseIndex[i],:] + K * (np.abs(self.X[SDchooseIndex[i],:]-self.X[-1,:])/(fitness[SDchooseIndex[i]] - fitness[-1] + 10E-8))

            self.BorderCheck()
            fitness = self.CaculateFitness()
            fitness,sortIndex = self.SortFitness(fitness)
            avg_fit.append(np.mean(fitness))
            avg_ts.append(np.var(fitness))
            self.X = self.SortPosition(sortIndex)
            if(fitness[0] < GbestScore):
                GbestScore = copy.copy(fitness[0])
                GbestPosition[0,:] = copy.copy(self.X[0,:])
            Curve[t+1] = GbestScore

        time_end = time.time()  # 记录迭代结束时间
        print(f'SSA2020共花费 {time_end - time_start} 秒')
        print('SSA2020最优适应度', GbestScore)
        print('SSA2020最优解', GbestPosition)
        if isBox:
            return Curve, self.X
        else:
            return Curve, avg_fit, avg_ts

